from torch import Tensor
import numpy as np
import os
import torch

def save_tensor_as_text(t: Tensor, filename: str):
    np_arr = np.array(t.detach())
    tmpfile = filename + '.tmp'
    np.savetxt(tmpfile, np_arr.flatten(), fmt="%0.d" if t.dtype == torch.int else "%.10f")
    with open(filename, 'w') as fp:
        fp.write(str(np_arr.shape)[1:-1] + "\n")
        fp.write(open(tmpfile).read())
    os.remove(tmpfile)

TEMPLATE_BASE ='''\
#include "dtype.h"

const [DATA_TYPE] [DATA_NAME][]={
[DATA_CONTENT]};

[TENSOR_TYPE] [TENSOR_NAME] = {
[TENSOR_CONTENT]};
'''

# t 为tensor型变量 name为字符串 eg. conv1.0.weight
# 本函数将其转换为C语言头文件
def save_tensor_as_c_header(t: Tensor, name: str, dir=''):
    length = len(t.shape)
    assert length in [1, 3, 4] # length=1: bias; length=3: featureMap; length=4: kernel;
    data_type = 'DTYPE'
    file_name = os.path.join(dir, name + '.h')
    tensor_name = name.replace('.', '_').replace('\\', '_').replace('//', '_')
    data_name = tensor_name + '_data'
    if length == 1: tensor_type = 'bias'
    if length == 3: tensor_type = 'featureMap'
    if length == 4: tensor_type = 'kernel'
    
    assert t.dtype in [torch.float32, torch.float64, torch.int16, torch.int32] 
    if t.dtype == torch.float32 or t.dtype == torch.float64:
        data_content = \
            list(map(lambda x: format(float(x), '.10f' if x>0 else '.9f') + ', ', t.flatten())) # 将tensor全部转化为浮点数字符串 长度截断10
    else:
        data_content = \
            list(map(lambda x: str(int(x)) + ', ', t.flatten())) # 将tensor全部转化为整数符串
        if 'bias' in name:
            data_type = 'long int'

    data_content = ['', *data_content]
    data_content[::8] = list(map(lambda x: x+'\n', data_content[::8]))
    data_content = ''.join(data_content[1:])
    
    tensor_content = ','.join(list(map(lambda x: str(int(x)), t.shape)))

    tensor_content += ',\n({}*)'.format(data_type) + data_name
    
    with open(file_name, 'w') as fp:
        fp.write(
        TEMPLATE_BASE
            .replace('[DATA_TYPE]', data_type)
            .replace('[DATA_NAME]', data_name)
            .replace('[DATA_CONTENT]', data_content)
            .replace('[TENSOR_TYPE]', tensor_type)
            .replace('[TENSOR_NAME]', tensor_name)
            .replace('[TENSOR_CONTENT]', tensor_content)
        )


def save_int16_tensor_as_c_header(t: Tensor, name: str):
    pass # TODO

def test():
    img = torch.randn(2, 10 , 10)
    save_tensor_as_c_header(img, 'img')

if __name__ == '__main__':
    test()